"""
FastAPI service for monitoring AgentMemory in real-time
Provides REST API endpoints and web interface for browser-based monitoring
"""

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, List, Optional, Any
import json
import asyncio
from datetime import datetime
import threading
import time

# AgentMemory monitoring service
class MemoryMonitorService:
    def __init__(self):
        self.agent_memory = None
        self.memory_history = []
        self.max_history_entries = 100
        self.start_time = datetime.now()
        self.is_running = False

    def set_agent_memory(self, agent_memory):
        """Set the AgentMemory instance to monitor"""
        self.agent_memory = agent_memory
        self.is_running = True

    def record_memory_state(self, event_type: str = "update"):
        """Record current memory state to history"""
        if not self.agent_memory:
            return

        state = {
            "timestamp": datetime.now().isoformat(),
            "event_type": event_type,
            "goal": self.agent_memory.goal,
            "current_step_id": self.agent_memory.current_step_id,
            "current_thought": self.agent_memory.current_thought,
            "current_action": self.agent_memory.current_action,
            "current_observation": str(self.agent_memory.current_observation)[:200] if self.agent_memory.current_observation else None,
            "iteration_count": self.agent_memory.iteration_count,
            "max_iteration": self.agent_memory.max_iteration,
            "error_message": self.agent_memory.error_message,
            "final_answer": self.agent_memory.final_answer,
            "history_length": len(self.agent_memory.history),
            "reasoning_steps_length": len(self.agent_memory.reasoning_steps)
        }

        self.memory_history.append(state)

        # Keep only recent history to prevent memory bloat
        if len(self.memory_history) > self.max_history_entries:
            self.memory_history = self.memory_history[-self.max_history_entries:]

    def get_current_state(self) -> Dict[str, Any]:
        """Get current memory state"""
        if not self.agent_memory:
            return {"error": "No AgentMemory instance available"}

        # Get context engine stats if available
        stats = {}
        if hasattr(self.agent_memory, 'context_engine'):
            try:
                stats = self.agent_memory.context_engine.get_stats()
            except:
                stats = {}

        return {
            "timestamp": datetime.now().isoformat(),
            "goal": self.agent_memory.goal,
            "current_step_id": self.agent_memory.current_step_id,
            "current_thought": self.agent_memory.current_thought,
            "current_action": self.agent_memory.current_action,
            "current_observation": self.agent_memory.current_observation,
            "iteration_count": self.agent_memory.iteration_count,
            "max_iteration": self.agent_memory.max_iteration,
            "error_message": self.agent_memory.error_message,
            "final_answer": self.agent_memory.final_answer,
            "history_length": len(self.agent_memory.history),
            "reasoning_steps_length": len(self.agent_memory.reasoning_steps),
            "context_stats": stats,
            "service_uptime": str(datetime.now() - self.start_time)
        }

    def get_detailed_history(self, limit: int = 10) -> List[Dict[str, Any]]:
        """Get detailed ReAct history"""
        if not self.agent_memory:
            return []

        history = []
        for i, (thought, action, observation) in enumerate(self.agent_memory.history[-limit:], 1):
            history.append({
                "step": i,
                "thought": thought,
                "action": action,
                "observation": str(observation)[:300] + "..." if len(str(observation)) > 300 else str(observation)
            })
        return history

    def get_insights(self) -> Dict[str, Any]:
        """Get context engine insights if available"""
        if not self.agent_memory or not hasattr(self.agent_memory, 'context_engine'):
            return {"error": "Context engine not available"}

        try:
            insights = self.agent_memory.context_engine.curator.insights
            formatted_insights = []
            for i, (trigger, insight) in enumerate(insights.items(), 1):
                formatted_insights.append({
                    "id": i,
                    "trigger": trigger,
                    "lesson": insight.lesson,
                    "weight": insight.weight,
                    "count": insight.count
                })
            return {"insights": formatted_insights, "count": len(formatted_insights)}
        except Exception as e:
            return {"error": f"Failed to get insights: {str(e)}"}

    def get_memory_timeline(self) -> List[Dict[str, Any]]:
        """Get memory state timeline"""
        return self.memory_history

# Global monitor instance
monitor_service = MemoryMonitorService()

# FastAPI app
app = FastAPI(
    title="AgentMemory Monitor",
    description="Real-time monitoring service for AgentMemory state",
    version="1.0.0"
)

# Enable CORS for web interface
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models for API responses
class MemoryState(BaseModel):
    timestamp: str
    goal: Optional[str]
    current_step_id: Optional[int]
    current_thought: Optional[str]
    current_action: Optional[Dict]
    current_observation: Optional[Any]
    iteration_count: int
    max_iteration: int
    error_message: Optional[str]
    final_answer: Optional[str]
    history_length: int
    reasoning_steps_length: int
    context_stats: Dict[str, Any]
    service_uptime: str

class HistoryEntry(BaseModel):
    step: int
    thought: str
    action: Dict
    observation: str

class InsightEntry(BaseModel):
    id: int
    trigger: str
    lesson: str
    weight: float
    count: int

# API Endpoints
@app.get("/", response_class=HTMLResponse)
async def get_web_interface():
    """Serve the web monitoring interface"""
    return HTMLResponse(content=get_monitor_html(), status_code=200)

@app.get("/api/state", response_model=MemoryState)
async def get_current_state():
    """Get current AgentMemory state"""
    state = monitor_service.get_current_state()
    if "error" in state:
        raise HTTPException(status_code=503, detail=state["error"])
    return state

@app.get("/api/history", response_model=List[HistoryEntry])
async def get_history(limit: int = 10):
    """Get ReAct history"""
    history = monitor_service.get_detailed_history(limit)
    return history

@app.get("/api/insights", response_model=Dict[str, Any])
async def get_insights():
    """Get context engine insights"""
    return monitor_service.get_insights()

@app.get("/api/timeline", response_model=List[Dict[str, Any]])
async def get_timeline():
    """Get memory state timeline"""
    return monitor_service.get_memory_timeline()

@app.post("/api/record/{event_type}")
async def record_event(event_type: str):
    """Record a memory state event"""
    monitor_service.record_memory_state(event_type)
    return {"status": "recorded", "event_type": event_type}

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service_running": monitor_service.is_running,
        "has_agent_memory": monitor_service.agent_memory is not None,
        "uptime": str(datetime.now() - monitor_service.start_time),
        "timestamp": datetime.now().isoformat()
    }

# Web interface HTML
def get_monitor_html() -> str:
    """Generate the monitoring web interface HTML"""
    return """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>AgentMemory Monitor</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }

        body {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }

        .container {
            max-width: 1200px;
            margin: 0 auto;
            background: white;
            border-radius: 15px;
            box-shadow: 0 20px 40px rgba(0,0,0,0.1);
            overflow: hidden;
        }

        .header {
            background: linear-gradient(135deg, #2c3e50 0%, #34495e 100%);
            color: white;
            padding: 30px;
            text-align: center;
        }

        .header h1 {
            font-size: 2.5em;
            margin-bottom: 10px;
            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        }

        .header p {
            font-size: 1.1em;
            opacity: 0.9;
        }

        .dashboard {
            padding: 30px;
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 30px;
        }

        .panel {
            background: #f8f9fa;
            border-radius: 10px;
            padding: 25px;
            border-left: 5px solid #3498db;
            box-shadow: 0 5px 15px rgba(0,0,0,0.08);
        }

        .panel h2 {
            color: #2c3e50;
            margin-bottom: 20px;
            font-size: 1.4em;
            display: flex;
            align-items: center;
        }

        .panel h2::before {
            content: "📊";
            margin-right: 10px;
            font-size: 1.2em;
        }

        .status-grid {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
            gap: 15px;
            margin-bottom: 20px;
        }

        .status-item {
            background: white;
            padding: 15px;
            border-radius: 8px;
            border: 1px solid #e9ecef;
            text-align: center;
        }

        .status-label {
            font-size: 0.9em;
            color: #6c757d;
            margin-bottom: 5px;
        }

        .status-value {
            font-size: 1.3em;
            font-weight: bold;
            color: #2c3e50;
        }

        .progress-bar {
            width: 100%;
            height: 8px;
            background: #e9ecef;
            border-radius: 4px;
            overflow: hidden;
            margin-top: 10px;
        }

        .progress-fill {
            height: 100%;
            background: linear-gradient(90deg, #3498db, #2ecc71);
            border-radius: 4px;
            transition: width 0.3s ease;
        }

        .current-goal {
            background: white;
            padding: 20px;
            border-radius: 8px;
            border: 1px solid #e9ecef;
            margin-bottom: 20px;
        }

        .current-thought {
            background: #e8f4f8;
            padding: 20px;
            border-radius: 8px;
            border-left: 4px solid #3498db;
            font-style: italic;
            margin-bottom: 20px;
        }

        .history-list {
            max-height: 300px;
            overflow-y: auto;
            background: white;
            border-radius: 8px;
            border: 1px solid #e9ecef;
        }

        .history-item {
            padding: 15px;
            border-bottom: 1px solid #f1f3f4;
        }

        .history-item:last-child {
            border-bottom: none;
        }

        .history-step {
            font-weight: bold;
            color: #3498db;
            margin-bottom: 8px;
        }

        .history-thought {
            color: #2c3e50;
            margin-bottom: 8px;
            font-style: italic;
        }

        .history-action {
            background: #f8f9fa;
            padding: 10px;
            border-radius: 5px;
            font-family: 'Courier New', monospace;
            font-size: 0.9em;
            color: #495057;
            margin-bottom: 8px;
        }

        .history-observation {
            color: #6c757d;
            font-size: 0.9em;
        }

        .insights-grid {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
            gap: 15px;
        }

        .insight-card {
            background: white;
            padding: 20px;
            border-radius: 8px;
            border: 1px solid #e9ecef;
            position: relative;
        }

        .insight-trigger {
            font-weight: bold;
            color: #e74c3c;
            margin-bottom: 10px;
        }

        .insight-lesson {
            color: #2c3e50;
            margin-bottom: 10px;
        }

        .insight-meta {
            display: flex;
            justify-content: space-between;
            font-size: 0.8em;
            color: #6c757d;
        }

        .refresh-button {
            background: linear-gradient(135deg, #3498db, #2980b9);
            color: white;
            border: none;
            padding: 12px 24px;
            border-radius: 25px;
            cursor: pointer;
            font-size: 1em;
            margin: 20px auto;
            display: block;
            transition: all 0.3s ease;
            box-shadow: 0 4px 15px rgba(52, 152, 219, 0.3);
        }

        .refresh-button:hover {
            transform: translateY(-2px);
            box-shadow: 0 6px 20px rgba(52, 152, 219, 0.4);
        }

        .error-message {
            background: #f8d7da;
            color: #721c24;
            padding: 15px;
            border-radius: 8px;
            border: 1px solid #f5c6cb;
            margin: 20px 0;
        }

        .loading {
            text-align: center;
            padding: 40px;
            color: #6c757d;
        }

        .full-width {
            grid-column: 1 / -1;
        }

        @media (max-width: 768px) {
            .dashboard {
                grid-template-columns: 1fr;
                gap: 20px;
                padding: 20px;
            }

            .header h1 {
                font-size: 2em;
            }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🧠 AgentMemory Monitor</h1>
            <p>Real-time monitoring of your AI agent's memory state and decision process</p>
        </div>

        <div class="dashboard">
            <div class="panel">
                <h2>Current Status</h2>
                <div id="current-status">
                    <div class="loading">Loading...</div>
                </div>
            </div>

            <div class="panel">
                <h2>Progress & Iterations</h2>
                <div id="progress-info">
                    <div class="loading">Loading...</div>
                </div>
            </div>

            <div class="panel full-width">
                <h2>Current Goal & Thought</h2>
                <div id="goal-thought">
                    <div class="loading">Loading...</div>
                </div>
            </div>

            <div class="panel">
                <h2>ReAct History</h2>
                <div id="react-history">
                    <div class="loading">Loading...</div>
                </div>
            </div>

            <div class="panel">
                <h2>Context Insights</h2>
                <div id="context-insights">
                    <div class="loading">Loading...</div>
                </div>
            </div>
        </div>

        <button class="refresh-button" onclick="refreshAllData()">🔄 Refresh Data</button>
    </div>

    <script>
        let refreshInterval;

        async function fetchData(endpoint) {
            try {
                const response = await fetch(`/api/${endpoint}`);
                if (!response.ok) {
                    throw new Error(`HTTP ${response.status}: ${response.statusText}`);
                }
                return await response.json();
            } catch (error) {
                console.error(`Error fetching ${endpoint}:`, error);
                return { error: error.message };
            }
        }

        function formatCurrentStatus(data) {
            if (data.error) {
                return `<div class="error-message">Error: ${data.error}</div>`;
            }

            return `
                <div class="status-grid">
                    <div class="status-item">
                        <div class="status-label">Goal Status</div>
                        <div class="status-value">${data.goal ? 'Active' : 'Idle'}</div>
                    </div>
                    <div class="status-item">
                        <div class="status-label">Current Step</div>
                        <div class="status-value">${data.current_step_id || 'None'}</div>
                    </div>
                    <div class="status-item">
                        <div class="status-label">History Length</div>
                        <div class="status-value">${data.history_length}</div>
                    </div>
                    <div class="status-item">
                        <div class="status-label">Service Uptime</div>
                        <div class="status-value">${data.service_uptime}</div>
                    </div>
                </div>
                ${data.error_message ? `<div class="error-message">${data.error_message}</div>` : ''}
                ${data.final_answer ? `<div class="current-goal"><strong>Final Answer:</strong><br>${data.final_answer}</div>` : ''}
            `;
        }

        function formatProgressInfo(data) {
            if (data.error) {
                return `<div class="error-message">Error: ${data.error}</div>`;
            }

            const progress = data.max_iteration > 0 ? (data.iteration_count / data.max_iteration) * 100 : 0;

            return `
                <div class="status-grid">
                    <div class="status-item">
                        <div class="status-label">Iterations</div>
                        <div class="status-value">${data.iteration_count} / ${data.max_iteration}</div>
                    </div>
                    <div class="status-item">
                        <div class="status-label">Progress</div>
                        <div class="status-value">${Math.round(progress)}%</div>
                    </div>
                </div>
                <div class="progress-bar">
                    <div class="progress-fill" style="width: ${progress}%"></div>
                </div>
                ${data.context_stats && Object.keys(data.context_stats).length > 0 ? `
                    <div class="status-grid" style="margin-top: 15px;">
                        <div class="status-item">
                            <div class="status-label">Tasks Processed</div>
                            <div class="status-value">${data.context_stats.tasks_processed || 0}</div>
                        </div>
                        <div class="status-item">
                            <div class="status-label">Insights</div>
                            <div class="status-value">${data.context_stats.total_insights || 0}</div>
                        </div>
                    </div>
                ` : ''}
            `;
        }

        function formatGoalThought(data) {
            if (data.error) {
                return `<div class="error-message">Error: ${data.error}</div>`;
            }

            return `
                ${data.goal ? `<div class="current-goal"><strong>Current Goal:</strong><br>${data.goal}</div>` : '<div class="current-goal">No active goal</div>'}
                ${data.current_thought ? `<div class="current-thought"><strong>Current Thought:</strong><br>${data.current_thought}</div>` : ''}
                ${data.current_action ? `<div class="current-goal"><strong>Current Action:</strong><br><code>${JSON.stringify(data.current_action, null, 2)}</code></div>` : ''}
            `;
        }

        function formatReactHistory(history) {
            if (!Array.isArray(history) || history.length === 0) {
                return '<div class="current-goal">No history available</div>';
            }

            return `
                <div class="history-list">
                    ${history.map(item => `
                        <div class="history-item">
                            <div class="history-step">Step ${item.step}</div>
                            <div class="history-thought">${item.thought}</div>
                            <div class="history-action">${JSON.stringify(item.action, null, 2)}</div>
                            <div class="history-observation">${item.observation}</div>
                        </div>
                    `).join('')}
                </div>
            `;
        }

        function formatInsights(data) {
            if (data.error) {
                return `<div class="error-message">Error: ${data.error}</div>`;
            }

            if (!data.insights || data.insights.length === 0) {
                return '<div class="current-goal">No insights available</div>';
            }

            return `
                <div class="insights-grid">
                    ${data.insights.map(insight => `
                        <div class="insight-card">
                            <div class="insight-trigger">${insight.trigger}</div>
                            <div class="insight-lesson">${insight.lesson}</div>
                            <div class="insight-meta">
                                <span>Weight: ${insight.weight.toFixed(1)}</span>
                                <span>Count: ${insight.count}</span>
                            </div>
                        </div>
                    `).join('')}
                </div>
            `;
        }

        async function refreshAllData() {
            // Fetch all data in parallel
            const [state, history, insights] = await Promise.all([
                fetchData('state'),
                fetchData('history'),
                fetchData('insights')
            ]);

            // Update UI
            document.getElementById('current-status').innerHTML = formatCurrentStatus(state);
            document.getElementById('progress-info').innerHTML = formatProgressInfo(state);
            document.getElementById('goal-thought').innerHTML = formatGoalThought(state);
            document.getElementById('react-history').innerHTML = formatReactHistory(history);
            document.getElementById('context-insights').innerHTML = formatInsights(insights);
        }

        // Initialize and start auto-refresh
        refreshAllData();
        refreshInterval = setInterval(refreshAllData, 2000); // Refresh every 2 seconds

        // Clean up on page unload
        window.addEventListener('beforeunload', () => {
            if (refreshInterval) {
                clearInterval(refreshInterval);
            }
        });
    </script>
</body>
</html>
    """

# Integration functions for ReActAgent
def start_monitor_server(agent_memory, host: str = "0.0.0.0", port: int = 8000):
    """Start the monitoring server with the given AgentMemory instance"""
    monitor_service.set_agent_memory(agent_memory)

    import uvicorn
    print(f"🌐 Starting AgentMemory Monitor on http://{host}:{port}")
    print(f"📊 Open your browser and navigate to http://localhost:{port} to view the monitoring interface")

    # Start server in a separate thread
    def run_server():
        uvicorn.run(app, host=host, port=port, log_level="info")

    server_thread = threading.Thread(target=run_server, daemon=True)
    server_thread.start()

    return server_thread

def record_memory_event(event_type: str = "update"):
    """Record a memory event (can be called from ReActAgent)"""
    monitor_service.record_memory_state(event_type)

if __name__ == "__main__":
    # Test the service
    print("🧪 Testing AgentMemory Monitor Service...")

    # Create a mock AgentMemory for testing
    class MockAgentMemory:
        def __init__(self):
            self.goal = "Test goal"
            self.current_step_id = 1
            self.current_thought = "Testing the monitoring service"
            self.current_action = {"name": "test_action", "args": {"param": "value"}}
            self.current_observation = "Test observation"
            self.iteration_count = 5
            self.max_iteration = 10
            self.error_message = None
            self.final_answer = None
            self.history = [("thought1", {"action": "act1"}, "obs1"), ("thought2", {"action": "act2"}, "obs2")]
            self.reasoning_steps = []

            class MockContextEngine:
                def get_stats(self):
                    return {
                        "tasks_processed": 3,
                        "total_insights": 2,
                        "memory_saved": "45%"
                    }

                class MockCurator:
                    insights = {
                        "error_pattern": type('Insight', (), {
                            'lesson': 'Always check file existence before reading',
                            'weight': 0.8,
                            'count': 3
                        })()
                    }

                curator = MockCurator()

            self.context_engine = MockContextEngine()

    mock_memory = MockAgentMemory()
    start_monitor_server(mock_memory, port=8001)

    # Keep the main thread alive
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n👋 Shutting down monitoring service...")